#pragma once

#include <iostream>
using namespace std;

template <class T>
struct RBTreeNode;

template <class T, class Ptr, class Ref>
struct __RBTreeIterator
{
    typedef RBTreeNode<T> rb_node;
    typedef __RBTreeIterator<T, Ptr, Ref> Self;
    typedef __RBTreeIterator<T, T *, T &> Iterator;
    typedef __RBTreeIterator<T, const T *, const T &> const_Iterator;

public:
    __RBTreeIterator(rb_node *node = nullptr)
        : _node(node)
    {
    }
    __RBTreeIterator(const Iterator &it)
        : _node(it._node)
    {
    }
    __RBTreeIterator(Self &self)
        :_node(self._node)
    {}

    Ptr operator->()
    {
        return &_node->_data;
    }

    Ref operator*()
    {
        return _node->_data;
    }

    Self &operator++()
    {
        if (_node->_right != nullptr)
        {
            rb_node *cur = _node->_right;
            while (cur->_left != nullptr)
            {
                cur = cur->_left;
            }
            _node = cur;
        }
        else
        {
            rb_node *parent = _node->_parent;
            rb_node *cur = _node;
            while (parent != nullptr && parent->_right == cur)
            {
                cur = parent;
                parent = parent->_parent;
            }

            if (parent != nullptr)
                _node = parent;
            else
                _node = nullptr;
        }
        return *this;
    }

    Self operator++(int)
    {
        rb_node *tmp = _node;
        if (_node->_right != nullptr)
        {
            rb_node *cur = _node->_right;
            while (cur->_left != nullptr)
            {
                cur = cur->_left;
            }
            _node = cur;
        }
        else
        {
            rb_node *parent = _node->_parent;
            rb_node *cur = _node;
            while (parent != nullptr && parent->_right == cur)
            {
                cur = parent;
                parent = parent->_parent;
            }

            if (parent != nullptr)
                _node = parent;
            else
                _node = nullptr;
        }
        return Self(tmp);
    }

    bool operator!=(const Self &self)
    {
        return _node != self._node;
    }

    Self &operator--()
    {
        if (_node->_left != nullptr)
        {
            rb_node *cur = _node->_left;

            while (cur->_right)
            {
                cur = cur->_right;
            }
            _node = cur;
        }
        else
        {
            rb_node *parent = _node->_parent;
            rb_node *cur = _node;
            while (parent != nullptr && parent->_left == cur)
            {
                cur = parent;
                parent = parent->_parent;
            }
            _node = parent;
        }
        return *this;
    }

    Self operator--(int)
    {
        rb_node *tmp = _node;
        if (_node->_left != nullptr)
        {
            rb_node *cur = _node->_left;

            while (cur->_right)
            {
                cur = cur->_right;
            }
            _node = cur;
        }
        else
        {
            rb_node *parent = _node->_parent;
            rb_node *cur = _node;
            while (parent != nullptr && parent->_left == cur)
            {
                cur = parent;
                parent = parent->_parent;
            }
            _node = parent;
        }
        return Self(tmp);
    }

public:
    rb_node *_node;
};

enum Color
{
    BLOCK,
    RED
};

template <class T>
struct RBTreeNode
{
    typedef RBTreeNode<T> _rb_node;

    RBTreeNode(const T &data = T())
        : _left(nullptr),
          _right(nullptr),
          _parent(nullptr),
          _col(RED),
          _data(data)
    {
    }

    _rb_node *_left;
    _rb_node *_right;
    _rb_node *_parent;
    T _data;
    Color _col;
};

template <class K, class T, class KeyOfT>
class RBTree
{
public:
    typedef RBTreeNode<T> rb_node;
    typedef __RBTreeIterator<T, T *, T &> Iterator;
    typedef __RBTreeIterator<T, const T *, const T &> const_Iterator;

    RBTree()
    {
        _root = nullptr;
    }

    ~RBTree()
    {
        Destroy(_root);
        _root = nullptr;
    }

    rb_node *Find(const K &k)
    {
        rb_node *cur = _root;
        while (cur)
        {
            if (key(_root->_data) < k)
            {
                cur = cur->_right;
            }
            else if (key(_root->_data) > k)
            {
                cur = cur->_left;
            }
            else
                break;
        }
        return cur;
    }

    Iterator Begin()
    {
        rb_node *cur = _root;
        while (cur->_left != nullptr)
        {
            cur = cur->_left;
        }
        return Iterator(cur);
    }

    Iterator End()
    {
        return Iterator(nullptr);
    }

    const_Iterator Begin() const
    {
        rb_node *cur = _root;
        while (cur->_left != nullptr)
        {
            cur = cur->_left;
        }
        return const_Iterator(cur);
    }

    const_Iterator End() const
    {
        return const_Iterator(nullptr);
    }

    pair<Iterator, bool> Insert(const T &data)
    {
        KeyOfT key; /*从data里面获取Key*/
        if (_root == nullptr)
        {
            _root = new rb_node(data);
            _root->_col = BLOCK;
            return make_pair(Iterator(_root), true);
        }
        rb_node *parent = nullptr;
        rb_node *cur = _root;
        while (cur)
        {
            if (key(cur->_data) > key(data)) /*当前值大于date的值， 那么data里面的key应该在左边*/
            {
                parent = cur;
                cur = cur->_left;
            }
            else if (key(cur->_data) < key(data))
            {
                parent = cur;
                cur = cur->_right;
            }
            else
            {
                return make_pair(Iterator(cur), false);
            }
        }
        /*查找到了末尾位置， 此时加入节点*/
        cur = new rb_node(data);
        rb_node *newnode = cur;

        if (key(parent->_data) > key(data))
            parent->_left = cur;
        else
            parent->_right = cur;

        cur->_parent = parent;

        /*开始向上更新*/
        while (parent && parent->_col == RED) /*父亲是红色， cur也是红色， 此时出现了连续的红色节点*/
        {
            rb_node *gparent = parent->_parent; /*创造一个祖父节点*/
            if (parent == gparent->_left)
            {
                //      g
                //   p     u
                rb_node *uncle = gparent->_right; /*叔叔节点是祖父节点的另一边*/
                if (uncle && uncle->_col == RED)
                {
                    /*变色*/
                    parent->_col = uncle->_col = BLOCK;
                    gparent->_col = RED;

                    /*向上处理*/
                    cur = gparent;
                    parent = cur->_parent;
                }
                else
                {
                    /*因为父亲节点是祖父节点的左边， cur是parent的左边， 所以就是单旋，类似AVL树*/
                    if (cur == parent->_left)
                    {
                        RotateR(gparent); /*对爷爷节点进行右单旋*/
                        parent->_col = BLOCK;
                        gparent->_col = RED;
                    }
                    else
                    {
                        /*双旋情况， parent是gparent的左， cur是parent的右*/
                        RotateL(parent);
                        RotateR(gparent);
                        cur->_col = BLOCK;
                        gparent->_col = parent->_col = RED;
                    }

                    /*旋转加变色完成之后， 这一次插入节点的操作就完毕了， 就可以break了。 */
                }
            }
            else
            {
                //      g
                //   u     p
                rb_node *uncle = gparent->_left;
                if (uncle && uncle->_col == RED)
                {
                    parent->_col = uncle->_col = BLOCK;
                    gparent->_col = RED;

                    cur = gparent;
                    parent = cur->_parent;
                }
                else /*旋转的情况*/
                {
                    if (cur == parent->_right)
                    {
                        RotateL(gparent);
                        parent->_col = BLOCK;
                        gparent->_col = RED;
                        cur->_col = RED;
                    }
                    else
                    {
                        RotateR(parent);
                        RotateL(gparent);
                        cur->_col = BLOCK;
                        gparent->_col = parent->_col = RED;
                    }

                    /*旋转完成之后， 此次插入节点操作就完成了， 就可以退出循环了*/
                    break;
                }
            }
        }

        _root->_col = BLOCK;
        return make_pair(Iterator(newnode), true);
    }

    rb_node *LeftMost()
    {
        rb_node *_ret = _root;
        while (_ret->_left)
            _ret = _ret->_left;
        return _ret;
    }

    rb_node *RightMost()
    {
        rb_node *_ret = _root;
        while (_ret->_right)
            _ret = _ret->_right;
        return _ret;
    }

    void InOrder()
    {
        _InOrder(_root);
    }

    bool IsBalance()
    {
        if (_root == nullptr)
            return false;
        if (_root->_col == RED)
            return false;

        int retNum = 0;
        rb_node *cur = _root;
        while (cur)
        {
            if (cur->_col == BLOCK)
            {
                retNum++;
            }
            cur = cur->_left;
        }
        return Check(_root, 0, retNum);
    }

    /**/

    bool Check(rb_node *root, int blackNum, const int retNum)
    {
        if (root == nullptr)
        {
            /*前序遍历走到空时，意味着一条路径走完了*/
            /**/
            if (retNum != blackNum)
            {
                cout << "存在黑色节点的数量不相等的路径" << endl;
                return false;
            }
            return true;
        }
        /*如果cur和父亲节点都是红色*/
        if (root->_col == RED && root->_parent->_col == RED)
        {
            cout << "存在连续的红色节点" << endl;
        }

        if (root->_col == BLOCK)
        {
            blackNum++; /*如果root是黑色， 那么让bloackNum加加即可， 到时候就比较retNum和blockNum的数量*/
        }
        return Check(root->_left, blackNum, retNum) &&
               Check(root->_right, blackNum, retNum);
    }

private:
    void Destroy(rb_node *root)
    {
        if (root == nullptr)
            return;
        Destroy(root->_left);
        Destroy(root->_right);
        delete root;
    }

    void _InOrder(rb_node *root)
    {
        if (root == nullptr)
            return;
        _InOrder(root->_left);
        cout << key(root->_data) << endl;
        _InOrder(root->_right);
    }

    void RotateL(rb_node *parent)
    {
        /*从上到下三次修改， 从下到上三次修改*/
        rb_node *gparent = parent->_parent;
        rb_node *subR = parent->_right;
        rb_node *subRL = subR->_left;

        if (gparent == nullptr)
            _root = subR;
        else
        {
            if (parent == gparent->_right)
                gparent->_right = subR;
            else
                gparent->_left = subR;
        }
        parent->_right = subRL;
        subR->_left = parent;

        if (subRL != nullptr)
            subRL->_parent = parent;
        parent->_parent = subR;
        subR->_parent = gparent;
    }

    void RotateR(rb_node *parent)
    {
        rb_node *gparent = parent->_parent;
        rb_node *subL = parent->_left;
        rb_node *subLR = subL->_right;

        if (gparent == nullptr)
            _root = subL;
        else
        {
            if (parent == gparent->_right)
                gparent->_right = subL;
            else
                gparent->_left = subL;
        }
        parent->_left = subLR;
        subL->_right = parent;

        if (subLR != nullptr)
            subLR->_parent = parent;
        parent->_parent = subL;
        subL->_parent = gparent;
    }

public:
    static void Test()
    {
    }

private:
    rb_node *_root;
    KeyOfT key;
};